#include "sms4.h"

#include <assert.h>
#include <stdlib.h>

#define LSL(x, n)       (((x) << ((n)%32)) + (((n)%32)==0 ? 0 : (x) >> (32-((n)%32))))

#define UINT8_TO_32(a0, a1, a2, a3) (((a0)<<24) + ((a1)<<16) + ((a2)<<8) + (a3))
#define UINT32_TO_8(i32, i80, i81, i82, i83) \
	((i80) = (uint8_t)((i32)>>24)), \
	((i81) = (uint8_t)((i32)>>16)), \
	((i82) = (uint8_t)((i32)>>8)), \
	((i83) = (uint8_t)(i32))

#define UINT32_TO_8S(i32, i8s) 	UINT32_TO_8(i32, (i8s)[0], (i8s)[1], (i8s)[2], (i8s)[3])
#define UINT8S_TO_32(i8s) 		UINT8_TO_32((i8s)[0], (i8s)[1], (i8s)[2], (i8s)[3])


// 2.2
static uint8_t Sbox(uint8_t d) {
	const static uint8_t S[16][16] = {
		{ 0xd6,0x90,0xe9,0xfe,0xcc,0xe1,0x3d,0xb7,0x16,0xb6,0x14,0xc2,0x28,0xfb,0x2c,0x05 },
		{ 0x2b,0x67,0x9a,0x76,0x2a,0xbe,0x04,0xc3,0xaa,0x44,0x13,0x26,0x49,0x86,0x06,0x99 },
		{ 0x9c,0x42,0x50,0xf4,0x91,0xef,0x98,0x7a,0x33,0x54,0x0b,0x43,0xed,0xcf,0xac,0x62 },
		{ 0xe4,0xb3,0x1c,0xa9,0xc9,0x08,0xe8,0x95,0x80,0xdf,0x94,0xfa,0x75,0x8f,0x3f,0xa6 },
		{ 0x47,0x07,0xa7,0xfc,0xf3,0x73,0x17,0xba,0x83,0x59,0x3c,0x19,0xe6,0x85,0x4f,0xa8 },
		{ 0x68,0x6b,0x81,0xb2,0x71,0x64,0xda,0x8b,0xf8,0xeb,0x0f,0x4b,0x70,0x56,0x9d,0x35 },
		{ 0x1e,0x24,0x0e,0x5e,0x63,0x58,0xd1,0xa2,0x25,0x22,0x7c,0x3b,0x01,0x21,0x78,0x87 },
		{ 0xd4,0x00,0x46,0x57,0x9f,0xd3,0x27,0x52,0x4c,0x36,0x02,0xe7,0xa0,0xc4,0xc8,0x9e },
		{ 0xea,0xbf,0x8a,0xd2,0x40,0xc7,0x38,0xb5,0xa3,0xf7,0xf2,0xce,0xf9,0x61,0x15,0xa1 },
		{ 0xe0,0xae,0x5d,0xa4,0x9b,0x34,0x1a,0x55,0xad,0x93,0x32,0x30,0xf5,0x8c,0xb1,0xe3 },
		{ 0x1d,0xf6,0xe2,0x2e,0x82,0x66,0xca,0x60,0xc0,0x29,0x23,0xab,0x0d,0x53,0x4e,0x6f },
		{ 0xd5,0xdb,0x37,0x45,0xde,0xfd,0x8e,0x2f,0x03,0xff,0x6a,0x72,0x6d,0x6c,0x5b,0x51 },
		{ 0x8d,0x1b,0xaf,0x92,0xbb,0xdd,0xbc,0x7f,0x11,0xd9,0x5c,0x41,0x1f,0x10,0x5a,0xd8 },
		{ 0x0a,0xc1,0x31,0x88,0xa5,0xcd,0x7b,0xbd,0x2d,0x74,0xd0,0x12,0xb8,0xe5,0xb4,0xb0 },
		{ 0x89,0x69,0x97,0x4a,0x0c,0x96,0x77,0x7e,0x65,0xb9,0xf1,0x09,0xc5,0x6e,0xc6,0x84 },
		{ 0x18,0xf0,0x7d,0xec,0x3a,0xdc,0x4d,0x20,0x79,0xee,0x5f,0x3e,0xd7,0xcb,0x39,0x48 }
	};

	return S[d>>4][d&0xF];
}

// 2.1.1
static uint32_t Tao(uint32_t d32) {
	uint8_t d8[4];
	UINT32_TO_8S(d32, d8);
	int a = Sbox(0xef);
	return UINT8_TO_32(Sbox(d8[0]), Sbox(d8[1]), Sbox(d8[2]), Sbox(d8[3]));
}


// 2.1.2
static uint32_t L(uint32_t b) {
	return b ^ LSL(b,2) ^ LSL(b, 10) ^ LSL(b, 18) ^ LSL(b, 24);
}

// 2.1
static uint32_t T(uint32_t d) {
	return L(Tao(d));
}

static uint32_t LPrime(uint32_t b) {
	return b ^ LSL(b, 13) ^ LSL(b, 23);
}

static uint32_t TPrime(uint32_t k) {
	return LPrime(Tao(k));
}

void SMS4_Initialize(SMS4_WorkArea *sms4, const uint8_t pw[SMS4_BLOCK_SIZE]) {
	assert(sms4 != NULL && pw != NULL);

	const static uint32_t fk[4] = {
		0xA3B1BAC6, 0x56AA3350, 0x677D9197, 0xB27022DC
	};

	const uint32_t ck[32] = { 
		0x00070e15, 0x1c232a31, 0x383f464d, 0x545b6269,
		0x70777e85, 0x8c939aa1, 0xa8afb6bd, 0xc4cbd2d9,
		0xe0e7eef5, 0xfc030a11, 0x181f262d, 0x343b4249,
		0x50575e65, 0x6c737a81, 0x888f969d, 0xa4abb2b9,
		0xc0c7ced5, 0xdce3eaf1, 0xf8ff060d, 0x141b2229,
		0x30373e45, 0x4c535a61, 0x686f767d, 0x848b9299,
		0xa0a7aeb5, 0xbcc3cad1, 0xd8dfe6ed, 0xf4fb0209,
		0x10171e25, 0x2c333a41, 0x484f565d, 0x646b7279
	};

	uint32_t mk[4] = { 
		UINT8S_TO_32(pw), UINT8S_TO_32(pw+4), 
		UINT8S_TO_32(pw+8), UINT8S_TO_32(pw+12)
	};

	uint32_t k[36];

	k[0] = mk[0] ^ fk[0];
	k[1] = mk[1] ^ fk[1];
	k[2] = mk[2] ^ fk[2];
	k[3] = mk[3] ^ fk[3];

	for (size_t i=0; i<32; ++i) {
		k[i+4] = k[i] ^ TPrime(k[i+1] ^ k[i+2] ^ k[i+3] ^ ck[i]);
		sms4->rk[i] = k[i+4];
	}
}

void SMS4_Encrypt(const SMS4_WorkArea *sms4, const uint8_t plain[SMS4_BLOCK_SIZE], uint8_t cipher[SMS4_BLOCK_SIZE])
{
	uint32_t x[36] = {
		UINT8S_TO_32(plain),   UINT8S_TO_32(plain+4),
		UINT8S_TO_32(plain+8), UINT8S_TO_32(plain+12)
	};

	for (int i=0; i<32; ++i) {
		x[i+4] =  x[i] ^ T(x[i+1] ^ x[i+2] ^ x[i+3] ^ sms4->rk[i]);
	}

	UINT32_TO_8S(x[35], cipher);
	UINT32_TO_8S(x[34], cipher + 4);
	UINT32_TO_8S(x[33], cipher + 8);
	UINT32_TO_8S(x[32], cipher + 12);

}

void SMS4_Decrypt(const SMS4_WorkArea *sms4, const uint8_t cipher[SMS4_BLOCK_SIZE], uint8_t plain[SMS4_BLOCK_SIZE]) {
	uint32_t y[36] = {
		UINT8S_TO_32(cipher),   UINT8S_TO_32(cipher+4),
		UINT8S_TO_32(cipher+8), UINT8S_TO_32(cipher+12)
	};

	for (int i=0; i<32; ++i) {
		y[i+4] =  y[i] ^ T(y[i+1] ^ y[i+2] ^ y[i+3] ^ sms4->rk[32-1-i]);
	}

	UINT32_TO_8S(y[35], plain);
	UINT32_TO_8S(y[34], plain + 4);
	UINT32_TO_8S(y[33], plain + 8);
	UINT32_TO_8S(y[32], plain + 12);

}

void SMS4_Finalize(SMS4_WorkArea *sms4) {
	for (int i=0; i<sizeof(sms4->rk)/sizeof(sms4->rk[0]); i++) {
		sms4->rk[i] = 0;
	}
}
